"""
大模型接口模块 - 负责与大语言模型进行交互

该模块提供与大语言模型通信的接口，将游戏状态转换为自然语言描述，
发送给大模型获取高级决策建议，并将建议转换为具体操作指令。
"""

import json
import requests
import os
from typing import Dict, Any, Optional, List
import numpy as np


class LLMInterface:
    """
    大语言模型接口类 - 与大语言模型进行交互
    
    该类负责：
    1. 将游戏状态转换为自然语言描述
    2. 与大语言模型通信获取决策建议
    3. 将决策建议转换为具体操作指令
    """
    
    def __init__(self, api_url: str = None, api_key: str = None):
        """
        初始化大模型接口
        
        Args:
            api_url: 大模型API地址（可选）
            api_key: API密钥（可选）
        """
        # 大模型API配置
        self.api_url = api_url or os.getenv("LLM_API_URL", "http://localhost:8000/v1/completions")
        self.api_key = api_key or os.getenv("LLM_API_KEY", "")
        
        # 游戏策略知识库
        self.strategy_knowledge = {
            "early_game": [
                "优先采集资源",
                "建造基础建筑",
                "侦察对手动向",
                "保持兵力平衡"
            ],
            "mid_game": [
                "扩张领土",
                "升级科技",
                "组建军队",
                "控制关键位置"
            ],
            "late_game": [
                "决战准备",
                "资源最大化",
                "高级兵种",
                "战术配合"
            ]
        }
        
        print("大模型接口初始化完成")
    
    def game_state_to_prompt(self, game_state: Dict[str, Any]) -> str:
        """
        将游戏状态转换为自然语言描述提示
        
        Args:
            game_state: 游戏状态字典
            
        Returns:
            str: 自然语言描述的提示文本
        """
        try:
            # 构建游戏状态描述
            prompt = "你是一个魔兽争霸3游戏专家AI助手。根据当前游戏状态，请给出最优策略建议。\n\n"
            prompt += "当前游戏状态：\n"
            
            # 资源信息
            resources = game_state.get('resources', {})
            prompt += f"- 金币: {resources.get('gold', 0)}\n"
            prompt += f"- 木材: {resources.get('wood', 0)}\n"
            
            # 单位和建筑信息
            prompt += f"- 单位数量: {game_state.get('unit_count', 0)}\n"
            prompt += f"- 建筑数量: {game_state.get('building_count', 0)}\n"
            
            # 游戏阶段判断（简化版）
            total_resources = resources.get('gold', 0) + resources.get('wood', 0)
            if total_resources < 1000:
                game_phase = "early_game"
                prompt += "- 游戏阶段: 前期\n"
            elif total_resources < 5000:
                game_phase = "mid_game"
                prompt += "- 游戏阶段: 中期\n"
            else:
                game_phase = "late_game"
                prompt += "- 游戏阶段: 后期\n"
            
            # 添加策略建议
            prompt += "\n推荐策略：\n"
            for strategy in self.strategy_knowledge.get(game_phase, []):
                prompt += f"- {strategy}\n"
            
            prompt += "\n请根据以上信息，给出具体的操作建议，包括：\n"
            prompt += "1. 当前最应该执行的操作类型（如采集、建造、训练、攻击等）\n"
            prompt += "2. 具体的操作目标或参数\n"
            prompt += "3. 操作的优先级和原因\n"
            prompt += "请以JSON格式返回结果，包含'action_type'、'target'、'priority'、'reason'字段。\n"
            
            return prompt
        except Exception as e:
            print(f"游戏状态转换失败: {e}")
            return ""
    
    def call_llm(self, prompt: str, max_tokens: int = 500) -> Optional[Dict[str, Any]]:
        """
        调用大语言模型API
        
        Args:
            prompt: 输入提示文本
            max_tokens: 最大生成token数
            
        Returns:
            Dict[str, Any]: 模型返回结果，失败返回None
        """
        try:
            # 构建请求头
            headers = {
                "Content-Type": "application/json"
            }
            
            # 如果有API密钥，则添加到请求头
            if self.api_key:
                headers["Authorization"] = f"Bearer {self.api_key}"
            
            # 构建请求数据
            data = {
                "prompt": prompt,
                "max_tokens": max_tokens,
                "temperature": 0.7,
                "stop": ["\n\n"]  # 停止标记
            }
            
            # 发送请求
            response = requests.post(self.api_url, headers=headers, json=data, timeout=30)
            
            # 检查响应状态
            if response.status_code == 200:
                result = response.json()
                return result
            else:
                print(f"大模型API调用失败，状态码: {response.status_code}")
                return None
        except Exception as e:
            print(f"大模型API调用异常: {e}")
            return None
    
    def parse_llm_response(self, response: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        解析大模型返回的响应
        
        Args:
            response: 大模型返回的响应
            
        Returns:
            Dict[str, Any]: 解析后的操作指令，失败返回None
        """
        try:
            # 提取生成的文本
            if 'choices' in response and len(response['choices']) > 0:
                text = response['choices'][0].get('text', '')
                
                # 尝试解析JSON格式的响应
                # 首先清理文本，提取JSON部分
                start_idx = text.find('{')
                end_idx = text.rfind('}') + 1
                
                if start_idx != -1 and end_idx > start_idx:
                    json_text = text[start_idx:end_idx]
                    parsed_data = json.loads(json_text)
                    
                    # 验证必要字段
                    required_fields = ['action_type', 'target', 'priority', 'reason']
                    if all(field in parsed_data for field in required_fields):
                        return parsed_data
                    else:
                        print("解析的JSON缺少必要字段")
                        return None
                else:
                    print("未找到有效的JSON格式响应")
                    return None
            else:
                print("响应中没有有效内容")
                return None
        except json.JSONDecodeError as e:
            print(f"JSON解析失败: {e}")
            return None
        except Exception as e:
            print(f"响应解析异常: {e}")
            return None
    
    def llm_decision(self, game_state: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        基于大模型的决策
        
        Args:
            game_state: 游戏状态
            
        Returns:
            Dict[str, Any]: 大模型建议的操作指令，失败返回None
        """
        try:
            # 将游戏状态转换为提示文本
            prompt = self.game_state_to_prompt(game_state)
            if not prompt:
                return None
            
            # 调用大模型
            response = self.call_llm(prompt)
            if not response:
                return None
            
            # 解析响应
            action = self.parse_llm_response(response)
            return action
        except Exception as e:
            print(f"大模型决策过程异常: {e}")
            return None
    
    def integrate_with_rl(self, rl_action: Dict[str, Any], 
                         llm_action: Optional[Dict[str, Any]]) -> Dict[str, Any]:
        """
        将大模型建议与强化学习决策进行融合
        
        Args:
            rl_action: 强化学习决策结果
            llm_action: 大模型建议的操作
            
        Returns:
            Dict[str, Any]: 融合后的最终决策
        """
        # 如果没有大模型建议，则直接返回强化学习决策
        if not llm_action:
            return rl_action
        
        try:
            # 简单的融合策略：根据优先级和置信度加权
            # 这里可以实现更复杂的融合算法
            
            # 示例：如果大模型建议的优先级很高，则采用大模型建议
            llm_priority = llm_action.get('priority', 0)
            if llm_priority > 8:  # 假设优先级满分是10
                # 转换大模型建议为标准操作格式，保持与RL决策一致
                integrated_decision = {
                    'source': 'llm',
                    'action': {
                        'type': llm_action.get('action_type', 'idle'),
                        'params': llm_action.get('target', {})
                    },
                    'confidence': llm_priority / 10.0,
                    'reason': llm_action.get('reason', ''),
                    'llm_raw': llm_action
                }
                return integrated_decision
            
            # 否则以强化学习决策为主
            integrated_decision = rl_action.copy()
            integrated_decision['source'] = 'rl_with_llm_advice'
            integrated_decision['llm_advice'] = llm_action
            
            return integrated_decision
        except Exception as e:
            print(f"决策融合过程异常: {e}")
            return rl_action


# 使用示例
if __name__ == "__main__":
    # 创建大模型接口实例
    llm_interface = LLMInterface()
    
    # 示例游戏状态
    sample_state = {
        'resources': {'gold': 800, 'wood': 450},
        'unit_count': 8,
        'building_count': 5
    }
    
    # 获取大模型建议（需要实际的大模型API支持）
    # llm_action = llm_interface.llm_decision(sample_state)
    # print("大模型建议:", llm_action)
    
    print("大模型接口模块测试完成")
